#include "matrix_util.hpp"
#include <gtest/gtest.h>

#define IGEMM_TEST(DIM_M, DIM_N, DIM_K, IB)                                    \
  TEST(TEST_IGEMM, M##DIM_M##_N##DIM_N##_K##DIM_K) {                           \
    unsigned long m = DIM_M;                                                   \
    unsigned long n = DIM_N;                                                   \
    unsigned long k = DIM_K;                                                   \
                                                                               \
    int *A = nullptr;                                                          \
    int *B = nullptr;                                                          \
    const std::size_t size_A = m * k;                                          \
    const std::size_t size_B = k * n;                                          \
    FIX_RAND_SEED(true);                                                       \
    ALLOC_INT_ARRAY_WITHOUT_INIT(A, size_A);                                   \
    ALLOC_INT_ARRAY_WITHOUT_INIT(B, size_B);                                   \
    RAND_ARRAY(A, size_A, IB);                                                 \
    RAND_ARRAY(B, size_B, IB);                                                 \
                                                                               \
    int *C = nullptr;                                                          \
    const std::size_t size_C = m * n;                                          \
    ALLOC_INT_ARRAY_WITHOUT_INIT(C, size_C);                                   \
    cblas_igemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0, A, k, \
                B, n, 0.0, C, n);                                              \
                                                                               \
    double *DA = nullptr;                                                      \
    double *DB = nullptr;                                                      \
    double *DC = nullptr;                                                      \
    ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DA, size_A);                               \
    ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DB, size_B);                               \
    ALLOC_DOUBLE_ARRAY_WITHOUT_INIT(DC, size_C);                               \
    COPY_ARRAY_I2D(A, DA, size_A);                                             \
    COPY_ARRAY_I2D(B, DB, size_B);                                             \
    cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, m, n, k, 1.0, DA,   \
                k, DB, n, 0.0, DC, n);                                         \
                                                                               \
    /* check answers  */                                                       \
    const long mp = 2;                                                         \
    const long np = 2;                                                         \
    const long mb = 512;                                                       \
    const long nb = 512;                                                       \
    for (std::size_t i = 0; i < mp; i++) {                                     \
      for (std::size_t j = 0; j < np; j++) {                                   \
        for (std::size_t ii = 0; ii < mb; ii++) {                              \
          for (std::size_t jj = 0; jj < nb; jj++) {                            \
            auto expect = DC[(i * mb + ii) * n + j * nb + jj];                 \
            auto actual = C[(i * mb + ii) * n + j * nb + jj];                  \
            ASSERT_EQ(expect, static_cast<double>(actual))                     \
                << "diff at (" << ii << ", " << jj << ") in BLK (" << i        \
                << ", " << j << ")";                                           \
          }                                                                    \
        }                                                                      \
      }                                                                        \
    }                                                                          \
    std::cout << DC[0] << " " << C[0] << std::endl;                            \
                                                                               \
    delete[] A;                                                                \
    delete[] B;                                                                \
    delete[] C;                                                                \
    delete[] DA;                                                               \
    delete[] DB;                                                               \
    delete[] DC;                                                               \
  }

IGEMM_TEST(1024, 1024, 1024, 7)